/*========================== begin_copyright_notice ============================

Copyright (C) 2022-2023 Intel Corporation

SPDX-License-Identifier: MIT

============================= end_copyright_notice ===========================*/

#include "AdaptorCommon/FreezeIntDiv.hpp"
#include "Compiler/IGCPassSupport.h"

#include "common/LLVMWarningsPush.hpp"
#include "llvm/IR/InstVisitor.h"
#include "llvm/IR/IRBuilder.h"
#include "common/LLVMWarningsPop.hpp"

using namespace llvm;

namespace IGC {

class FreezeIntDiv : public FunctionPass, public InstVisitor<FreezeIntDiv> {
public:
  static char ID;

  FreezeIntDiv() : FunctionPass(ID), changed(false) { initializeFreezeIntDivPass(*PassRegistry::getPassRegistry()); }

  bool runOnFunction(Function &F) override;

  StringRef getPassName() const override { return "FreezeIntDiv"; }

  void getAnalysisUsage(AnalysisUsage &AU) const override { AU.setPreservesCFG(); }

  void visitBinaryOperator(BinaryOperator &I);
  void freezeIntDiv(BinaryOperator &I);

private:
  bool changed;
};

char FreezeIntDiv::ID = 0;

#define PASS_FLAG "igc-freeze-int-div-pass"
#define PASS_DESCRIPTION "Freeze integer division"
#define PASS_CFG_ONLY false
#define PASS_ANALYSIS false
IGC_INITIALIZE_PASS_BEGIN(FreezeIntDiv, PASS_FLAG, PASS_DESCRIPTION, PASS_CFG_ONLY, PASS_ANALYSIS)
IGC_INITIALIZE_PASS_END(FreezeIntDiv, PASS_FLAG, PASS_DESCRIPTION, PASS_CFG_ONLY, PASS_ANALYSIS)

FunctionPass *createFreezeIntDivPass() { return new FreezeIntDiv(); }

////////////////////////////////////////////////////////////////////////////
bool FreezeIntDiv::runOnFunction(Function &F) {
  changed = false;
  visit(F);
  return changed;
}

void FreezeIntDiv::freezeIntDiv(BinaryOperator &I) {
  llvm::IRBuilder<> builder(I.getNextNode());
  Value *FI = builder.CreateFreeze(&I, "freeze");
  I.replaceAllUsesWith(FI);
  cast<Instruction>(FI)->setOperand(0, &I);
}

void FreezeIntDiv::visitBinaryOperator(BinaryOperator &I) {
  switch (I.getOpcode()) {
  case Instruction::UDiv:
  case Instruction::SDiv:
  case Instruction::URem:
  case Instruction::SRem:
    // Stop propagation of poison values generated by LLVM in case of
    // integer division by zero. Use LLVM 10+ freeze instruction.
    {
      // In theory, we shouldn't have to check whether the input is, in
      // fact, poison/undef - the resulting freeze should simply resolve
      // to a no-op otherwise. However, until we resolve all freeze insts
      // in codegen, LLVM passes (InstCombine, GVN, SimplifyCFG, etc.)
      // can be sloppy with checking the operand's validity - as a
      // result, valuable optimizations can be skipped in some cases.
      // Therefore, minimize freeze emission for cases that are
      // guaranteed to be valid.
      auto *constDivisor = dyn_cast<ConstantInt>(I.getOperand(1));
      if (constDivisor && !constDivisor->isZero())
        break;
      freezeIntDiv(I);
      changed = true;
      break;
    }
  default:
    break;
  }
}

} // namespace IGC
